import h5py
import os
import sys
import argparse
import logging

def parse_arguments():
    """解析命令行参数。"""
    parser = argparse.ArgumentParser(description="修改 HDF5 文件中数据集的 chunksize 和压缩设置，并保留所有组和文件级别的属性。")
    parser.add_argument('--case_name', required=True, help='案例名称，例如 "SPE10"')
    parser.add_argument('--output_path', required=True, help='输出目录，例如 "output/"')
    parser.add_argument('--new_chunks', type=int, default=262144, help='新的 chunksize，默认为 262144。如果设置为 0，则使用连续存储并禁用压缩。')
    parser.add_argument('--compression', type=str, default=None, help='新的压缩类型，例如 "gzip", "lzf"。仅在使用 chunked 存储时有效。')
    parser.add_argument('--compression_opts', type=int, default=None, help='压缩选项，例如 gzip 的压缩级别（1-9）。')
    return parser.parse_args()

def setup_logging():
    """设置日志配置。"""
    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')

def copy_attributes(src, dst):
    """复制 HDF5 对象的所有属性。"""
    for key, value in src.attrs.items():
        dst.attrs[key] = value

def copy_items(name, obj, src, dst, new_chunks, new_compression, new_compression_opts):
    """
    递归复制 HDF5 对象，并根据需要修改数据集的 chunksize 和压缩设置。

    参数：
    - name: HDF5 对象的名称。
    - obj: HDF5 对象。
    - src: 源 HDF5 文件对象。
    - dst: 目标 HDF5 文件对象。
    - new_chunks: 新的 chunksize 值。如果为 0，则使用连续存储并禁用压缩。
    - new_compression: 新的压缩类型（字符串）或 None。
    - new_compression_opts: 新的压缩选项（整数）或 None。
    """
    if isinstance(obj, h5py.Group):
        try:
            dst_group = dst.create_group(name)
            copy_attributes(obj, dst_group)
            logging.info(f"复制组: {name}")
        except Exception as e:
            logging.error(f"无法复制组 {name}: {e}")
    elif isinstance(obj, h5py.Dataset):
        try:
            logging.info(f"复制数据集: {name}")
            logging.debug(f"  Shape: {obj.shape}")
            logging.debug(f"  原始 Chunksize: {obj.chunks}")
            logging.debug(f"  Compression: {obj.compression}")
            logging.debug(f"  Compression opts: {obj.compression_opts}")
            logging.debug(f"  Shuffle: {obj.shuffle}")
            logging.debug(f"  Fletcher32: {obj.fletcher32}")
            logging.debug(f"  Maxshape: {obj.maxshape}")
            logging.debug(f"  Dtype: {obj.dtype}")

            dtype = obj.dtype
            shape = obj.shape
            compression = obj.compression
            compression_opts = obj.compression_opts
            shuffle = obj.shuffle
            fletcher32 = obj.fletcher32
            maxshape = obj.maxshape

            # 确定新的 chunksize 和压缩设置
            if new_chunks > 0 and len(shape) > 0:
                # 使用指定的 chunksize
                if new_chunks <= shape[0]:
                    chunks = (new_chunks,) + shape[1:]
                else:
                    chunks = obj.chunks
                # 使用指定的压缩设置
                if new_compression is not None:
                    compression = new_compression
                if new_compression_opts is not None:
                    compression_opts = new_compression_opts
            else:
                # 使用连续存储，禁用压缩和可扩展性
                chunks = None
                compression = None
                compression_opts = None
                shuffle = False
                fletcher32 = False
                maxshape = shape

            # 准备创建数据集的参数
            dset_kwargs = {
                'shape': shape,
                'dtype': dtype,
                'compression': compression,
                'compression_opts': compression_opts,
                'shuffle': shuffle,
                'fletcher32': fletcher32,
                'maxshape': maxshape
            }

            if chunks is not None:
                dset_kwargs['chunks'] = chunks

            # 打印 dset_kwargs 以进行调试
            logging.debug(f"dset_kwargs for {name}: {dset_kwargs}")

            # 在目标文件中创建新的数据集
            dst_dset = dst.create_dataset(name, **dset_kwargs)

            # 复制数据
            # if len(shape) == 0:
                # # 标量数据集
                # dst_dset[...] = obj[...]
            # else:
                # 多维数据集，考虑分块复制以节省内存
                # block_size = 1024  # 可根据需要调整块大小
                # if shape[0] > block_size:
                    # for i in range(0, shape[0], block_size):
                        # end = min(i + block_size, shape[0])
                        # dst_dset[i:end] = obj[i:end]
                # else:
                    # dst_dset[...] = obj[...]
            dst_dset[...] = obj[...]

            # 复制数据集的属性
            copy_attributes(obj, dst_dset)
            logging.info(f"成功复制数据集: {name}")

        except Exception as e:
            logging.error(f"无法复制数据集 {name}: {e}")

def main():
    """主函数，执行 HDF5 文件的复制和修改。"""
    args = parse_arguments()
    setup_logging()

    output_path = args.output_path
    case_name = args.case_name
    original_file = f"{case_name}_timeseries.vtkhdf"
    new_file = f"{case_name}_modified.vtkhdf"

    input_file = os.path.join(output_path, case_name, original_file)
    output_file = os.path.join(output_path, case_name, new_file)

    new_chunks = args.new_chunks
    new_compression = args.compression
    new_compression_opts = args.compression_opts

    logging.info(f"输入文件: {input_file}")
    logging.info(f"输出文件: {output_file}")
    if new_chunks > 0:
        logging.info(f"新的 chunksize: {new_chunks}")
        if new_compression:
            logging.info(f"新的压缩类型: {new_compression}")
        if new_compression_opts is not None:
            logging.info(f"新的压缩选项: {new_compression_opts}")
    else:
        logging.info("使用连续存储 (chunksize=None) 并禁用压缩")

    # 检查输入文件是否存在
    if not os.path.isfile(input_file):
        logging.error(f"输入文件 '{input_file}' 不存在。")
        sys.exit(1)

    # 创建输出目录（如果不存在）
    os.makedirs(os.path.dirname(output_file), exist_ok=True)

    with h5py.File(input_file, 'r') as src, h5py.File(output_file, 'w') as dst:
        # 遍历源文件中的所有对象并复制
        src.visititems(lambda name, obj: copy_items(name, obj, src, dst, new_chunks, new_compression, new_compression_opts))

        # 复制文件级别的属性
        copy_attributes(src, dst)
        logging.info("复制文件级别的属性完成。")

    logging.info(f"Chunksize 和压缩设置已成功修改，新的文件为: {new_file}")

if __name__ == "__main__":
    main()
